import torch
import matplotlib.pyplot as plt
import BuildingTheNetWorking as B
import Training as T
import torch.optim as optim
import SettingUp as s
import Test as ts

continued_network = B.Net()
continued_optimizer = optim.SGD(B.network.parameters(), lr=s.learning_rate, momentum=s.momentum)

network_state_dict = torch.load('model.pth')
continued_network.load_state_dict(network_state_dict)

optimizer_state_dict = torch.load('optimizer.pth')
continued_optimizer.load_state_dict(optimizer_state_dict)

for i in range(4,9):
  T.test_counter.append(i*len(s.train_loader.dataset))
  T.train(i)
  ts.test()

fig = plt.figure()
plt.plot(T.train_counter, T.train_losses, color='blue')
plt.scatter(T.test_counter, T.test_losses, color='red')
plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likelihood loss')
plt.show()